import warnings
warnings.filterwarnings("ignore")

import os
# 使用 Keras 3 的多后端能力：将后端切换为 PyTorch
os.environ.setdefault("KERAS_BACKEND", "torch")

from keras import layers, models
from keras.datasets import imdb
from keras.utils import pad_sequences
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score
import seaborn as sns
import matplotlib.pyplot as plt

def build_model(vocab_size: int, embedding_dim: int, lstm_units: int, max_len: int):
    model = models.Sequential([
        layers.Embedding(input_dim=vocab_size, output_dim=embedding_dim, input_length=max_len),
        layers.LSTM(lstm_units, dropout=0.2, recurrent_dropout=0.2),
        layers.Dense(1, activation='sigmoid')
    ])
    model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
    return model

def main():
    print("开始加载 IMDB 情感数据集...")
    vocab_size = 20000
    max_len = 200
    embedding_dim = 128
    lstm_units = 128
    batch_size = 64
    epochs = 5

    (x_train, y_train), (x_test, y_test) = imdb.load_data(num_words=vocab_size)
    print("训练集大小:", len(x_train), "测试集大小:", len(x_test))

    print("进行序列填充/截断...")
    x_train = pad_sequences(x_train, maxlen=max_len)
    x_test = pad_sequences(x_test, maxlen=max_len)

    print("构建并训练 LSTM 模型...")
    model = build_model(vocab_size, embedding_dim, lstm_units, max_len)
    history = model.fit(
        x_train, y_train,
        validation_split=0.2,
        batch_size=batch_size,
        epochs=epochs,
        verbose=1
    )

    print("在测试集上评估...")
    y_prob = model.predict(x_test, batch_size=batch_size, verbose=0).ravel()
    y_pred = (y_prob >= 0.5).astype('int32')

    print("测试集准确率:", accuracy_score(y_test, y_pred))
    print("分类报告:\n")
    print(classification_report(y_test, y_pred, target_names=["负面", "正面"]))

    print("绘制混淆矩阵...")
    cm = confusion_matrix(y_test, y_pred)
    plt.figure(figsize=(6, 5))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', cbar=True)
    plt.title("混淆矩阵（IMDB 情感分析 LSTM）")
    plt.xlabel("预测类别")
    plt.ylabel("真实类别")
    plt.tight_layout()
    plt.show()

if __name__ == '__main__':
    main()


